#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import os
import re
from collections import OrderedDict

import click
from tqdm import tqdm


class GenomicInterval:
    __slots__ = [
        "chrom", "start",
        "end", "strand"
    ]

    def __init__(self, chrom: str, start: int, end: int, strand: str):
        u"""

        """
        self.chrom = chrom
        self.start = int(start)
        self.end = int(end)
        self.strand = strand

        if self.end < self.start:
            raise ValueError("{} should bigger than {}".format(self.end, self.start))

    def __hash__(self):
        return hash((self.chrom, self.start, self.end, self.strand))

    def __str__(self):
        return "{}:{}-{}:{}".format(self.chrom, self.start, self.end, self.strand)

    def is_overlapped(self, other):
        return self.chrom == other.chrom and \
               self.end > other.start and \
               self.start < other.end and \
               self.strand == other.strand

    def __eq__(self, other):
        return self.chrom == other.chrom and \
               self.start == other.start and \
               self.end == other.end and \
               self.strand == other.strand

    def __gt__(self, other):
        if self.chrom != other.chrom:
            return self.chrom > other.chrom

        if self.start != other.start:
            return self.start > other.start

        if self.end != other.end:
            return self.end > other.end

        return self.strand > other.strand

    def __lt__(self, other):
        if self.chrom != other.chrom:
            return self.chrom < other.chrom

        if self.start != other.start:
            return self.start < other.start

        if self.end != other.end:
            return self.end < other.end

        return self.strand < other.strand

    def to_dict(self) -> dict:
        return {
            "chrom": self.chrom,
            "start": self.start,
            "end": self.end,
            "strand": self.strand
        }


class GenomicFeature:
    __slots__ = [
        "name", "type", "iv",
        "source", "frame", "attr"
    ]

    def __init__(self, name: str, type_: str, interval: GenomicInterval, source: str = "."):
        u"""

        """

        self.name = name
        self.type = type_
        self.iv = interval
        self.source = source
        self.frame = "."
        self.attr = OrderedDict({})

    def __str__(self):
        u"""
        Convert GenomicFeature to gtf
        """
        attr_str = ['{} "{}";'.format(k, v) for k, v in self.attr.items()]

        return "{chrom}\t{source}\t{type}\t{start}\t{end}\t{frame}\t{strand}\t.\t{attr}".format(
            **{
                "source": self.source,
                "type": self.type,
                "frame": self.frame,
                "attr": " ".join(attr_str),
                **self.iv.to_dict(),
            }
        )

    def is_overlapped(self, other) -> bool:
        return self.iv.is_overlapped(other.iv)

    def cut(self, other) -> list:
        data = sorted([self, other], key=lambda x: (x.iv.chrom, x.iv.start, x.iv.end, x.iv.strand))

        res = []

        try:
            res.append(GenomicFeature(
                name=data[0].name,
                type_="exonic_part",
                interval=GenomicInterval(
                    chrom=data[0].iv.chrom,
                    start=data[0].iv.start,
                    end=data[1].iv.start - 1,
                    strand=data[0].iv.strand
                ),
                source=data[0].source
            ))
        except ValueError as _:
            pass

        try:
            res.append(GenomicFeature(
                name="+".join(sorted([data[0].name, data[1].name])),
                type_="exonic_part",
                interval=GenomicInterval(
                    chrom=data[0].iv.chrom,
                    start=data[1].iv.start,
                    end=min(data[1].iv.end, data[0].iv.end),
                    strand=data[0].iv.strand
                ),
                source=data[0].source
            ))
        except ValueError as _:
            pass

        try:
            res.append(GenomicFeature(
                name=data[0].name if data[1].iv.end < data[0].iv.end else data[1].name,
                type_="exonic_part",
                interval=GenomicInterval(
                    chrom=data[0].iv.chrom,
                    start=min(data[1].iv.end, data[0].iv.end) + 1,
                    end=max(data[1].iv.end, data[0].iv.end),
                    strand=data[0].iv.strand
                ),
                source=data[0].source
            ))
        except ValueError as _:
            pass

        return res

    def __eq__(self, other):
        return self.iv == other.iv

    def __gt__(self, other):
        return self.iv > other.iv

    def __lt__(self, other):
        return self.iv < other.iv

    def __hash__(self):
        return self.iv.__hash__()

    def to_exonic_part(self):
        return GenomicFeature(
            name=self.name,
            type_="exonic_part",
            interval=self.iv,
            source=self.source
        )

    def to_exonic_part_str(self) -> str:
        attr_str = ['{}:{}'.format(k, v) for k, v in self.attr.items()]

        return "{chrom}\t{source}\t{type}\t{start}\t{end}\t{frame}\t{strand}\t.\t{attr}".format(
            **{
                "source": self.source,
                "type": self.type,
                "frame": self.frame,
                "attr": " ".join(attr_str),
                **self.iv.to_dict(),
            }
        )


def __decode_attr__(info: str) -> dict:
    """
    decode attr into OrderedDict
    """
    res = OrderedDict({})

    if "=" in info:
        for line in info.split(";"):
            line = line.strip().split("=")

            res[line[0]] = line[-1]

    else:
        for line in info.split(";"):

            line = line.strip().split()
            if len(line) < 2:
                continue
            res[line[0]] = line[-1].replace('"', '')

    return res


def gff_reader(path: str, only_exon: bool = False):
    u"""
    generator of GenomicFeature
    :param path path to input gtf or gff3 file
    :param only_exon: only generate exons
    """
    # two dict to collect relationships between different elements
    genes_trans = {}
    trans_others = {}

    bar = tqdm(total=os.path.getsize(path), desc="Reading")

    read = 0
    with open(path) as r:

        while read < os.path.getsize(path):
            line = r.readline()
            bar.update(r.tell() - read)
            read = r.tell()

            if line.startswith("#"):
                continue

            lines = line.split()

            if len(lines) < 8:
                continue

            attrs = __decode_attr__(" ".join(lines[8:]))

            # format attrs
            if re.search(r"(transcript|rna)", lines[2], re.I):
                genes_trans[attrs.get("ID", attrs.get("transcript_id"))] = attrs.get("Parent", attrs.get("gene_id"))

                if "gene_id" not in attrs.keys():
                    attrs["gene_id"] = attrs.pop("Parent")

            elif lines[2] == "exon":
                trans_others[attrs.get("ID", attrs.get("Parent"))] = attrs.get("Parent", attrs.get("transcript_id"))

                if "gene_id" not in attrs.keys():
                    temp.attr["gene_id"] = genes_trans[attrs["Parent"]]
                    temp.attr["transcript_id"] = attrs.pop("Parent")

            if only_exon and lines[2] != "exon":
                continue

            try:
                temp_iv = GenomicInterval(
                    chrom=lines[0],
                    start=int(lines[3]),
                    end=int(lines[4]),
                    strand=lines[6]
                )

                temp = GenomicFeature(
                    name=attrs["gene_id"] if lines[2] != "exon" else attrs["transcript_id"],
                    type_=lines[2],
                    source=lines[1],
                    interval=temp_iv
                )

                for k, v in attrs.items():
                    temp.attr[k] = v

                yield temp
            except ValueError as err:
                print(line.strip(), err)

    bar.close()


def exonic_parts(exons: list) -> list:
    u""""
    function to make list of exons into parts
    :param list of GenomicFeatures
    :param list of partial GenomicFeatures
    """

    res = []
    exons = sorted(exons)

    if len(exons) == 1:
        res.append(exons[0].to_exonic_part())
    else:
        curr = exons[0]
        temp = []
        for next_ in exons[1:]:
            if curr.is_overlapped(next_):
                temp += curr.cut(next_)
                curr = temp.pop(-1)
            else:
                if curr not in temp:
                    temp.append(curr.to_exonic_part())

                res += temp
                curr = next_
                temp = []

        if len(temp) > 0:
            res += temp

    # assign id
    curr = res[0]
    curr.attr["001"] = curr.name
    known_id = set(curr.name.split("+"))

    curr_id = 2
    for i in res[1:]:
        temp_id = set(i.name.split("+"))

        if len(known_id | temp_id) < len(known_id) + len(temp_id):
            i.attr["%03d" % curr_id] = i.name
            known_id |= temp_id
            curr_id += 1
        else:
            known_id = temp_id
            i.attr["001"] = i.name
            curr_id = 2

    return res


def get_exonic_parts(infile: str, outfile: str):
    u"""

    """
    exons = {}
    for i in gff_reader(infile, only_exon=True):
        key = "{}#{}".format(i.iv.chrom, i.iv.strand)

        temp = exons.get(key, [])
        temp.append(i)
        exons[key] = temp

    res = []
    for values in tqdm(exons.values(), desc="Part"):
        res += exonic_parts(values)

    with open(outfile, "w+") as w:
        for data in tqdm(sorted(res), desc="Write"):
            w.write("{}\n".format(data.to_exonic_part_str()))


@click.command(
    context_settings=dict(help_option_names=['-h', '--help'])
)
@click.option(
    "-i",
    "--infile",
    type=click.Path(exists=True),
    required=True,
    help="Path to input gtf"
)
@click.option(
    "-o",
    "--outfile",
    type=click.Path(),
    required=True,
    help="Path to input gtf"
)
def main(infile, outfile):
    u"""
    Welcome
    """
    get_exonic_parts(infile, outfile)


if __name__ == '__main__':
    main()
